{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-08-07T07:08:14.805635Z",
     "start_time": "2024-08-07T07:08:12.737406Z"
    }
   },
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings\n",
    "from llama_index.embeddings.ollama import OllamaEmbedding\n",
    "from llama_index.llms.ollama import Ollama"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-07T07:09:27.535824Z",
     "start_time": "2024-08-07T07:08:14.806641Z"
    }
   },
   "cell_type": "code",
   "source": [
    "documents = SimpleDirectoryReader(\"data\").load_data()\n",
    "\n",
    "# nomic embedding model\n",
    "Settings.embed_model = OllamaEmbedding(model_name=\"nomic-embed-text\")\n",
    "\n",
    "# ollama\n",
    "Settings.llm = Ollama(model=\"llama3.1\", request_timeout=360.0)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    ")"
   ],
   "id": "8522879d2ecdeef0",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-07T07:10:31.935669Z",
     "start_time": "2024-08-07T07:09:27.535824Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 查询数据\n",
    "query_engine = index.as_query_engine().\n",
    "response = query_engine.query(\"Datawhale是什么?\")\n",
    "print(response)"
   ],
   "id": "d3ce752f58ed051f",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datawhale是一个专注于数据科学与 AI 领域的开源组织。\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-07T07:11:36.617356Z",
     "start_time": "2024-08-07T07:10:31.936734Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 检索上下文进行对话\n",
    "from llama_index.core.memory import ChatMemoryBuffer\n",
    "memory = ChatMemoryBuffer.from_defaults(token_limit=1500)\n",
    "\n",
    "chat_engine = index.as_chat_engine(\n",
    "    chat_mode=\"context\",\n",
    "    memory=memory,\n",
    "    system_prompt=(\n",
    "        \"You are a chatbot, able to have normal interactions.\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "response = chat_engine.chat(\"Datawhale是什么？\")\n",
    "print(response)"
   ],
   "id": "d54fbd73a6836902",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datawhale是一个开源组织，成立于2018年。它是一个聚集了众多领域院校和知名企业优秀学习者的团队，致力于构建AI领域的开源学习社区。\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-07T07:12:37.150782Z",
     "start_time": "2024-08-07T07:11:36.617356Z"
    }
   },
   "cell_type": "code",
   "source": [
    "response = chat_engine.chat(\"Datawhale是哪一年成立的？\")\n",
    "print(response)"
   ],
   "id": "aabf12effd3f3d9b",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datawhale成立于2018年。\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-07T07:14:32.060279Z",
     "start_time": "2024-08-07T07:13:15.546428Z"
    }
   },
   "cell_type": "code",
   "source": [
    "response = chat_engine.chat(\"Datawhale目标是什么？\")\n",
    "print(response)"
   ],
   "id": "8184df223ef9ca64",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Datawhale的目标是为学习者提供一个平台，让他们能与人一起成长、真实地展现自我、开放包容、互信互助、敢于试错和勇于担当，促进AI领域的开源学习和人才培养。\n"
     ]
    }
   ],
   "execution_count": 10
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
